import pandas as pd
import re


def clean_data(input_file, output_file):
    # 加载数据并处理潜在的结构问题
    df = pd.read_csv(input_file)

    # 移除重复的表头和全空行
    df = df[~df.apply(lambda x: x.astype(str).str.contains('学校名称|学校')).any(axis=1)]
    df = df.dropna(how='all')

    # 重置索引
    df = df.reset_index(drop=True)

    # 标准化列名
    df = df.rename(columns={
        '序号': 'serial_no',
        '学校名称': 'school_name',
        '学校': 'school_name',
        '学校性质': 'school_type',
        '教职工人数': 'total_staff_count',
        '专职教师人数': 'full_time_teacher_count',
        '师生比': 'student_faculty_ratio_raw'
    })

    # 筛选公办学校
    df = df[df['school_type'] == '公办']

    # 清理师生比列
    def extract_ratio(ratio_str):
        if pd.notna(ratio_str):
            match = re.search(r'1:(\d+\.?\d*)', str(ratio_str))
            if match:
                return float(match.group(1))
        return None

    df['student_faculty_ratio'] = df['student_faculty_ratio_raw'].apply(extract_ratio)

    # 计算估算学生总数
    df['estimated_student_count'] = df['total_staff_count'] * df['student_faculty_ratio']

    # 计算每位专职教师对应的学生数，处理除零错误
    def calculate_students_per_ft_teacher(row):
        if row['full_time_teacher_count'] == 0:
            return None
        return row['estimated_student_count'] / row['full_time_teacher_count']

    df['students_per_ft_teacher'] = df.apply(calculate_students_per_ft_teacher, axis=1)

    # 确保数值类型
    numeric_columns = ['total_staff_count', 'full_time_teacher_count',
                       'estimated_student_count', 'students_per_ft_teacher']
    for col in numeric_columns:
        df[col] = pd.to_numeric(df[col], errors='coerce')

    # 处理缺失数据
    df = df.dropna(subset=numeric_columns)

    # 选择并排序最终列
    final_columns = [
        'serial_no', 'school_name', 'total_staff_count',
        'full_time_teacher_count', 'estimated_student_count',
        'students_per_ft_teacher'
    ]
    df = df[final_columns]

    # 保存输出
    df.to_csv(output_file, index=False)
    print(f"已成功保存清理后的数据到 {output_file}")


if __name__ == "__main__":
    input_file = "baoshan-schools-2020.csv"
    output_file = "clean-baoshan-schools-2020.csv"
    clean_data(input_file, output_file)